use linfa::prelude::*;
use linfa::traits::Fit;
use linfa_linear::LinearRegression;
use ndarray::{Array1, Array2, ArrayView1, ArrayView2, Axis, concatenate};
use std::collections::VecDeque;
use std::f64::consts::PI;
use chrono::{NaiveDate, Duration};
use rand::seq::SliceRandom;
use rand::thread_rng;

// 数据结构：存储报废事件
#[derive(Debug, Clone)]
struct ScrapEvent {
    timestamp: f64, // 时间戳（从参考日期开始的天数）
    scrap_amount: f64, // 该时间点的报废数量
}

// 预测系统配置
#[derive(Clone)]
struct ForecastConfig {
    max_history: usize,    // 最大历史数据点数
    min_history: usize,    // 最小历史数据点数
    train_interval: f64,   // 训练间隔（天）
    validation_split: f64, // 验证集比例
    period: f64,           // 季节性周期（天）
    fourier_terms: usize,  // 傅里叶项数
    min_interval: f64,     // 最小时间间隔阈值（天）
    regularization: f64,   // 正则化强度
}

// 预测系统状态
struct ForecastingSystem {
    config: ForecastConfig,
    events: VecDeque<ScrapEvent>,       // 报废事件记录
    cumulative_scrap: Vec<DataPoint>,   // 累计报废量序列
    model: Option<LinearRegression<f64>>, // 线性回归模型
    last_train_time: f64,               // 最后训练时间
    reference_date: NaiveDate,          // 参考日期
    validation_error: Option<f64>,      // 验证集误差
}

// 累计报废量数据点
#[derive(Debug, Clone)]
struct DataPoint {
    timestamp: f64,
    cumulative_scrap: f64,
}

// 特征类型
#[derive(Debug, Clone)]
struct FeaturePoint {
    timestamp: f64,
    scrap_rate: f64,
    weight: f64,
}

impl ForecastConfig {
    fn default() -> Self {
        Self {
            max_history: 36,
            min_history: 3,
            train_interval: 30.0,
            validation_split: 0.3, // 30% 数据用于验证
            period: 365.0,
            fourier_terms: 2,
            min_interval: 0.1,
            regularization: 0.1, // L2 正则化强度
        }
    }
}

impl ForecastingSystem {
    fn new(reference_date: NaiveDate) -> Self {
        let config = ForecastConfig::default();
        Self {
            config,
            events: VecDeque::with_capacity(config.max_history),
            cumulative_scrap: Vec::new(),
            model: None,
            last_train_time: 0.0,
            reference_date,
            validation_error: None,
        }
    }

    // 添加新报废事件
    fn add_scrap_event(&mut self, date: NaiveDate, scrap_amount: f64) {
        // 计算时间戳（从参考日期开始的天数）
        let days_since_ref = (date - self.reference_date).num_days() as f64;
        
        // 确保时间顺序
        if let Some(last_event) = self.events.back() {
            if days_since_ref < last_event.timestamp {
                eprintln!("错误：新事件日期 {} 早于最后事件日期", date);
                return;
            }
        }
        
        // 添加事件记录
        if self.events.len() >= self.config.max_history {
            self.events.pop_front();
        }
        self.events.push_back(ScrapEvent {
            timestamp: days_since_ref,
            scrap_amount,
        });
        
        // 重建累计报废序列
        self.rebuild_cumulative_series();
        
        // 检查是否需要重新训练
        if self.should_retrain(days_since_ref) {
            self.train_model();
        }
    }
    
    // 重建累计报废序列
    fn rebuild_cumulative_series(&mut self) {
        self.cumulative_scrap.clear();
        let mut cumulative = 0.0;
        
        // 添加初始点（参考日期，报废量0）
        self.cumulative_scrap.push(DataPoint {
            timestamp: 0.0,
            cumulative_scrap: 0.0,
        });
        
        // 按时间排序事件
        let mut sorted_events: Vec<_> = self.events.iter().collect();
        sorted_events.sort_by(|a, b| a.timestamp.partial_cmp(&b.timestamp).unwrap());
        
        for event in sorted_events {
            cumulative += event.scrap_amount;
            self.cumulative_scrap.push(DataPoint {
                timestamp: event.timestamp,
                cumulative_scrap: cumulative,
            });
        }
    }

    // 判断是否需要重新训练
    fn should_retrain(&self, current_time: f64) -> bool {
        if self.cumulative_scrap.len() < self.config.min_history + 1 {
            return false;
        }
        current_time - self.last_train_time >= self.config.train_interval
    }

    // 准备特征数据
    fn prepare_features(&self) -> Vec<FeaturePoint> {
        let mut features = Vec::new();
        let current_time = self.cumulative_scrap.last().map(|p| p.timestamp).unwrap_or(0.0);
        
        for i in 1..self.cumulative_scrap.len() {
            let prev = &self.cumulative_scrap[i-1];
            let curr = &self.cumulative_scrap[i];
            
            let interval = curr.timestamp - prev.timestamp;
            if interval < self.config.min_interval {
                continue;
            }
            
            let scrap_amount = curr.cumulative_scrap - prev.cumulative_scrap;
            if scrap_amount <= 0.0 {
                continue;
            }
            
            let scrap_rate = scrap_amount / interval;
            let midpoint_time = (prev.timestamp + curr.timestamp) / 2.0;
            
            // 计算权重（基于时间衰减和间隔长度）
            let time_weight = (-0.01 * (current_time - curr.timestamp)).exp(); // 每日衰减1%
            let interval_weight = interval.sqrt();
            let weight = time_weight * interval_weight;
            
            features.push(FeaturePoint {
                timestamp: midpoint_time,
                scrap_rate,
                weight,
            });
        }
        
        features
    }

    // 构建傅里叶特征矩阵
    fn build_fourier_features(&self, times: &[f64]) -> Array2<f64> {
        let n = times.len();
        let total_features = 2 + 2 * self.config.fourier_terms; // 常数项 + 时间项 + 傅里叶项
        
        let mut features = Array2::zeros((n, total_features));
        
        for (i, &t) in times.iter().enumerate() {
            // 常数项
            features[[i, 0]] = 1.0;
            
            // 时间项
            features[[i, 1]] = t;
            
            // 傅里叶项
            for k in 0..self.config.fourier_terms {
                let k_f64 = (k + 1) as f64;
                let angle = 2.0 * PI * k_f64 * t / self.config.period;
                features[[i, 2 + 2 * k]] = angle.sin();
                features[[i, 3 + 2 * k]] = angle.cos();
            }
        }
        
        features
    }

    // 训练模型（带训练集/验证集分割）
    fn train_model(&mut self) {
        let features = self.prepare_features();
        
        if features.len() < self.config.min_history {
            println!("训练数据不足: 需要 {} 个点，实际有 {} 个", 
                self.config.min_history, features.len());
            return;
        }
        
        // 打乱数据并分割
        let mut rng = thread_rng();
        let mut shuffled_features = features.clone();
        shuffled_features.shuffle(&mut rng);
        
        let split_index = (shuffled_features.len() as f64 * (1.0 - self.config.validation_split)) as usize;
        let (train_features, valid_features) = shuffled_features.split_at(split_index);
        
        // 准备训练数据
        let train_times: Vec<f64> = train_features.iter().map(|fp| fp.timestamp).collect();
        let train_rates: Vec<f64> = train_features.iter().map(|fp| fp.scrap_rate).collect();
        let train_weights: Vec<f64> = train_features.iter().map(|fp| fp.weight).collect();
        
        // 准备验证数据
        let valid_times: Vec<f64> = valid_features.iter().map(|fp| fp.timestamp).collect();
        let valid_rates: Vec<f64> = valid_features.iter().map(|fp| fp.scrap_rate).collect();
        
        // 构建特征矩阵
        let train_x = self.build_fourier_features(&train_times);
        let valid_x = self.build_fourier_features(&valid_times);
        
        // 应用权重
        let train_y = Array1::from(train_rates);
        let weights = Array1::from(train_weights).mapv(|w| w.sqrt()); // 用于加权最小二乘
        
        // 加权特征矩阵和目标向量
        let weighted_x = train_x.clone() * weights.clone().insert_axis(Axis(1));
        let weighted_y = train_y.clone() * &weights;
        
        // 添加正则化项
        let mut regularized_x = weighted_x.clone();
        for i in 2..weighted_x.ncols() { // 跳过常数项和时间项
            regularized_x[[i, i]] += self.config.regularization;
        }
        
        // 创建数据集
        let dataset = Dataset::new(weighted_x.to_owned(), weighted_y.to_owned());
        
        // 训练线性回归模型
        let model = LinearRegression::new()
            .with_solver(linfa_linear::LinearRegressionParams::qr())
            .fit(&dataset)
            .expect("模型训练失败");
        
        // 在验证集上评估
        let valid_pred = model.predict(&valid_x);
        let valid_rates_arr = Array1::from(valid_rates);
        
        // 计算平均绝对百分比误差
        let mut mape = 0.0;
        let mut count = 0;
        for i in 0..valid_rates_arr.len() {
            let actual = valid_rates_arr[i].max(1e-6); // 避免除零
            let error = (valid_pred[i] - actual).abs() / actual;
            mape += error;
            count += 1;
        }
        
        if count > 0 {
            self.validation_error = Some(mape / count as f64);
        }
        
        println!("模型训练成功！训练集大小: {}, 验证集大小: {}, MAPE: {:.1}%", 
            train_features.len(), valid_features.len(), 
            self.validation_error.unwrap_or(0.0) * 100.0);
        
        self.model = Some(model);
        self.last_train_time = self.cumulative_scrap.last().unwrap().timestamp;
    }

    // 预测未来时间点的报废率
    fn predict_scrap_rate(&self, time: f64) -> f64 {
        if let Some(model) = &self.model {
            // 构建特征向量
            let mut features = Array1::zeros(2 + 2 * self.config.fourier_terms);
            features[0] = 1.0; // 常数项
            features[1] = time; // 时间项
            
            // 傅里叶项
            for k in 0..self.config.fourier_terms {
                let k_f64 = (k + 1) as f64;
                let angle = 2.0 * PI * k_f64 * time / self.config.period;
                features[2 + 2 * k] = angle.sin();
                features[3 + 2 * k] = angle.cos();
            }
            
            model.predict(&features.insert_axis(Axis(0)))[0].max(0.0)
        } else {
            // 没有模型时使用历史平均值
            self.calculate_historical_average()
        }
    }

    // 预测库存归零时间（天）
    fn predict_inventory_depletion(&self, current_inventory: f64, current_date: NaiveDate) -> Option<f64> {
        if current_inventory <= 0.0 {
            return Some(0.0);
        }
        
        // 计算当前时间戳
        let current_time = (current_date - self.reference_date).num_days() as f64;
        
        // 获取历史平均报废率
        let historical_rate = self.calculate_historical_average();
        
        // 使用模型预测当前报废率
        let current_rate = self.predict_scrap_rate(current_time).max(0.1); // 最小0.1避免除零
        
        // 简单估算：剩余库存 / 当前报废率
        let mut time_to_depletion = current_inventory / current_rate;
        
        // 使用模型进行精确积分预测
        let mut remaining = current_inventory;
        let mut time = current_time;
        let step = 0.1; // 时间步长（天）
        let max_time = current_time + 2.0 * time_to_depletion.max(365.0); // 最大预测时间（最多2年）
        
        while remaining > 0.0 && time < max_time {
            let rate = self.predict_scrap_rate(time).max(0.1);
            let consumption = rate * step;
            
            if consumption >= remaining {
                let fraction = remaining / consumption;
                time_to_depletion = (time + fraction * step) - current_time;
                break;
            }
            
            remaining -= consumption;
            time += step;
        }
        
        Some(time_to_depletion)
    }
    
    // 计算历史平均报废率
    fn calculate_historical_average(&self) -> f64 {
        if self.cumulative_scrap.len() < 2 {
            return 1.0; // 默认值
        }
        
        let mut total_scrap = 0.0;
        let mut total_days = 0.0;
        
        for i in 1..self.cumulative_scrap.len() {
            let prev = &self.cumulative_scrap[i-1];
            let curr = &self.cumulative_scrap[i];
            
            let interval = curr.timestamp - prev.timestamp;
            if interval < self.config.min_interval {
                continue;
            }
            
            let scrap_amount = curr.cumulative_scrap - prev.cumulative_scrap;
            if scrap_amount > 0.0 {
                total_scrap += scrap_amount;
                total_days += interval;
            }
        }
        
        if total_days > 0.0 {
            total_scrap / total_days
        } else {
            // 没有有效区间时使用总平均值
            let total_scrap = self.cumulative_scrap.last().unwrap().cumulative_scrap;
            let total_days = self.cumulative_scrap.last().unwrap().timestamp;
            
            if total_days > 0.0 {
                total_scrap / total_days
            } else {
                1.0
            }
        }
    }
    
    // 获取总报废量
    fn total_scrap(&self) -> f64 {
        self.cumulative_scrap.last().map_or(0.0, |p| p.cumulative_scrap)
    }
    
    // 获取验证误差
    fn validation_error(&self) -> Option<f64> {
        self.validation_error
    }
}

// 示例使用
fn main() {
    // 设置参考日期（最早事件的日期）
    let reference_date = NaiveDate::from_ymd_opt(2024, 1, 1).unwrap();
    let mut system = ForecastingSystem::new(reference_date);
    
    // 添加报废事件（日期，报废数量）
    system.add_scrap_event(NaiveDate::from_ymd_opt(2024, 1, 1).unwrap(), 32.0);
    system.add_scrap_event(NaiveDate::from_ymd_opt(2024, 2, 2).unwrap(), 60.0);
    system.add_scrap_event(NaiveDate::from_ymd_opt(2024, 4, 5).unwrap(), 100.0);
    system.add_scrap_event(NaiveDate::from_ymd_opt(2024, 6, 15).unwrap(), 85.0);
    system.add_scrap_event(NaiveDate::from_ymd_opt(2024, 8, 20).unwrap(), 110.0);
    system.add_scrap_event(NaiveDate::from_ymd_opt(2024, 10, 10).unwrap(), 95.0);
    system.add_scrap_event(NaiveDate::from_ymd_opt(2024, 12, 25).unwrap(), 78.0);
    system.add_scrap_event(NaiveDate::from_ymd_opt(2025, 2, 14).unwrap(), 105.0);
    system.add_scrap_event(NaiveDate::from_ymd_opt(2025, 4, 30).unwrap(), 90.0);
    
    println!("已添加 {} 个报废事件，总报废量: {}", 
        system.events.len(), system.total_scrap());
    
    // 用户输入当前库存和日期
    let current_inventory = 500.0;
    let current_date = NaiveDate::from_ymd_opt(2025, 10, 24).unwrap();
    
    // 预测库存归零时间（天）
    if let Some(days_to_depletion) = system.predict_inventory_depletion(current_inventory, current_date) {
        let depletion_date = current_date + Duration::days(days_to_depletion as i64);
        println!("\n当前库存 {} 件将在 {:.1} 天后耗尽", current_inventory, days_to_depletion);
        println!("预计耗尽日期: {}", depletion_date);
        
        // 预测未来报废率
        println!("\n未来报废率预测:");
        for days in [30, 60, 90] {
            let future_date = current_date + Duration::days(days);
            let future_time = (future_date - reference_date).num_days() as f64;
            let rate = system.predict_scrap_rate(future_time);
            println!("  {}天后 ({}) 预测报废率: {:.1}件/天", days, future_date, rate);
        }
        
        // 显示模型性能
        if let Some(mape) = system.validation_error() {
            println!("\n模型验证集平均绝对百分比误差: {:.1}%", mape * 100.0);
        }
    } else {
        println!("\n无法预测库存耗尽时间");
    }
}